# Copyright (c) 2024 Microsoft Corporation.
# Licensed under the MIT License

"""All the steps to transform final entities."""

from uuid import uuid4

import pandas as pd

from graphrag.config.models.embed_graph_config import EmbedGraphConfig
from graphrag.data_model.schemas import ENTITIES_FINAL_COLUMNS
from graphrag.index.operations.compute_degree import compute_degree
from graphrag.index.operations.create_graph import create_graph
from graphrag.index.operations.embed_graph.embed_graph import embed_graph
from graphrag.index.operations.layout_graph.layout_graph import layout_graph


def finalize_entities(
    entities: pd.DataFrame,
    relationships: pd.DataFrame,
    embed_config: EmbedGraphConfig | None = None,
    layout_enabled: bool = False,    # 是否启用 UMAP 布局计算
) -> pd.DataFrame:
    """All the steps to transform final entities."""

    # 步骤 1: 从关系数据构建 NetworkX 图,保留边的权重属性
    # 这个图用于后续的度数计算、嵌入生成和布局计算
    graph = create_graph(relationships, edge_attr=["weight"])

    # 步骤 2: 可选的图嵌入生成
    # 初始化图嵌入为 None
    graph_embeddings = None

    # 如果提供了嵌入配置且已启用,则使用 Node2Vec 算法生成图嵌入
    # 图嵌入可用于后续的 UMAP 降维,提供更好的可视化效果
    if embed_config is not None and embed_config.enabled:
        graph_embeddings = embed_graph(
            graph,
            embed_config,
        )

    # 步骤 3: 计算节点的 x/y 坐标用于可视化
    # 如果 layout_enabled=True 且有图嵌入,使用 UMAP 降维生成 2D 坐标
    # 如果 layout_enabled=False,所有节点将获得默认的 (0, 0) 坐标
    layout = layout_graph(
        graph,
        layout_enabled,
        embeddings=graph_embeddings,
    )

    # 步骤 4: 计算每个节点的度数(连接数)
    # 度数表示该实体在图中与多少其他实体有关系
    degrees = compute_degree(graph)


    # 步骤 5: 合并数据
    # 将布局信息(x, y 坐标)和度数信息合并到实体 DataFrame 中
    # 使用左连接确保保留所有原始实体,即使它们没有布局或度数信息
    final_entities = (
        entities.merge(layout, left_on="title", right_on="label", how="left")
        .merge(degrees, on="title", how="left")
        .drop_duplicates(subset="title")
    )

    # 步骤 6: 过滤和重置索引
    # 只保留 title 不为空的实体,并重置索引
    final_entities = final_entities.loc[entities["title"].notna()].reset_index()
    # disconnected nodes and those with no community even at level 0 can be missing degree

    # 步骤 7: 填充缺失的度数值
    # 断开连接的节点(孤立节点)和在第 0 层都没有社区的节点可能缺少度数
    # 将这些缺失值填充为 0,并转换为整数类型
    final_entities["degree"] = final_entities["degree"].fillna(0).astype(int)

    # 步骤 8: 再次重置索引以获得连续的索引值
    # 这个索引将用作 human_readable_id
    final_entities.reset_index(inplace=True)

    # 步骤 9: 分配两种类型的 ID
    # human_readable_id: 连续的整数索引,便于人类阅读和调试
    final_entities["human_readable_id"] = final_entities.index

    # id: 唯一的 UUID 字符串,用于系统内部的唯一标识
    # 使用 lambda 函数为每个实体生成一个新的 UUID
    final_entities["id"] = final_entities["human_readable_id"].apply(
        lambda _x: str(uuid4())
    )

    # 步骤 10: 返回最终结果
    # 只返回 ENTITIES_FINAL_COLUMNS 中定义的列,确保输出格式一致
    return final_entities.loc[
        :,
        ENTITIES_FINAL_COLUMNS,
    ]
